-
Notifications
You must be signed in to change notification settings - Fork 19.6k
Model Export to liteRT #21674
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Model Export to liteRT #21674
Conversation
Introduces a custom LiteRTExporter for exporting models to TFLite format, bypassing the standard TFLiteConverter. Updates the export API and documentation to support the new 'lite_rt' format, and adds relevant options for custom ops, select TF ops, and optimizations.
Replaces the custom MLIR-based TFLite conversion logic in LiteRTExporter with direct use of the standard TFLiteConverter. Also improves input signature handling for tf.function tracing and updates imports accordingly.
Moved imports of get_input_signature and make_tf_tensor_spec inside functions in saved_model.py to prevent circular imports. Updated EXPORT_FORMATS in export_utils.py to use string references instead of direct imports.
Adds max_sequence_length parameter to input signature generation for sequence models, bounding sequence length for transformer-like architectures. Improves LiteRTExporter with heuristics for complex models, fallback conversion via SavedModel for large models, and exposes max_sequence_length in Model export options. Updates documentation accordingly.
Adds logic to dynamically reduce max_sequence_length for large vocabulary models in export_utils.py to prevent tensor size overflow. In lite_rt_exporter.py, introduces checks and workarounds for models with _DictWrapper issues, and applies memory optimizations for large models during TFLite conversion. These changes improve export reliability and prevent memory errors for models such as Gemma, Llama, and similar architectures.
Removed custom trackable object logic from LiteRTExporter and now save the model directly, simplifying the export process. Also streamlined vocabulary size checks in export_utils to prevent tensor size overflow, removing verbose warnings and redundant comments.
Refactors the TFLite conversion logic in lite_rt_exporter.py to attempt direct conversion first and only fall back to SavedModel if necessary, improving robustness and clarity. Adds a new lite_rt_exporter_simple.py file with a streamlined LiteRTExporter class for direct TFLite export, bypassing complex MLIR conversion paths.
Refactors export_utils and lite_rt_exporter to better detect large vocabulary and Keras-Hub models, applying safer sequence length limits and more robust TFLite conversion paths. Adds heuristics for model type detection, ensures memory safety, and improves handling of TensorFlow introspection issues during export.
Working well with keras
Eliminates the logic for bounding sequence length in model export utilities and related code paths. The max_sequence_length parameter and associated shape bounding for large vocabulary models are removed from export_utils.py and lite_rt_exporter.py. Updates model export documentation accordingly. Adds a comprehensive test script for Keras Hub LiteRT export, verifying numerical accuracy between original and exported models.
Renames the exporter module and class from 'lite_rt_exporter' and 'LiteRTExporter' to 'litert_exporter' and 'LitertExporter', respectively. Updates all references, registry keys, and documentation to use 'litert' instead of 'lite_rt'. Adds the new 'export_litert' function and corresponding tests for Litert export functionality.
Refactors and improves the LiteRT exporter for better input signature inference, error handling, and verbose output. Updates and expands the test suite to cover multi-input, multi-output, and error scenarios, and ensures the Model.export() API is used consistently. Also improves docstrings and formatting for clarity and maintainability.
Replaces direct imports and availability checks for LiteRT with the LazyModule-based `litert` utility from `module_utils`. Updates tests and exporter logic to require '.tflite' filepaths and consistently use the new import pattern, improving modularity and error handling.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request introduces support for exporting Keras models to the LiteRT (TFLite) format, which is a valuable addition for mobile and edge deployment. The changes are well-structured, including a new LitertExporter
, updates to the Model.export
API and documentation, and related refactoring of export utilities. The implementation appears robust, with thoughtful handling of different model types and conversion strategies, and is accompanied by a comprehensive test suite. My feedback focuses on a few areas to improve code clarity, maintainability, and adherence to best practices.
Hello @pctablet505! Thanks for reaching out. I'm here to assist with your pull request, "Model Export to liteRT". I noticed from the Codecov report that the new changes have low patch coverage (10.20408%). This is primarily due to missing coverage in If you have any questions about improving the test coverage, clarifying specific parts of the code, or anything else related to this PR, please let me know how I can help! |
Removed the return value and conditional messaging from export_litert. The function now always prints a simple message upon export and no longer returns the compilation result or filepath.
Simplifies and standardizes input signature generation for Functional, Sequential, and subclassed models to always use a flexible batch dimension (None) for export. Refactors Sequential input shape inference to remove hardcoded fallbacks and ensure dynamic batching support. Cleans up and clarifies error handling and structure of input signature utilities.
Renamed litert_exporter.py and its test to litert.py and litert_test.py, respectively. Updated all relevant imports to use the new module name. Refactored LitertExporter to simplify model building logic and removed redundant input shape inference for Sequential models. Also cleaned up export_utils.py by removing the export format registry and related functions.
Moved imports of get_input_signature and make_tf_tensor_spec to the module level in saved_model.py for consistency and efficiency. In litert.py, imports remain local to preserve lazy loading where needed.
Hello @pctablet505! Thanks for reaching out again. I'm still here to help with your pull request, "Model Export to liteRT". My previous comment highlighted the low patch coverage (12.73585%) in Would you like me to provide some specific suggestions on how to improve the test coverage for these files, or do you have any other questions about the PR that I can assist with? |
LiteRT export and related tests now import TensorFlow conditionally, allowing Keras to function without TensorFlow installed. If TensorFlow is unavailable, LiteRT export is disabled and tests are skipped, with informative messages and error handling added to guide users.
This pull request adds support for exporting Keras models to the LiteRT (TFLite) format, along with improvements to input signature handling and export utility documentation. The changes ensure that LiteRT export is only available when TensorFlow is installed, update the
export
API and documentation, and enhance input signature inference for various model types.LiteRT Export Support:
LitertExporter
andexport_litert
inkeras/src/export/__init__.py
, making LiteRT export available only if TensorFlow is installed.Model.export
method to support the"litert"
format, including new options for LiteRT export and user-facing documentation and example. Raises an informative error if TensorFlow is not installed. [1] [2] [3] [4]litert
as a lazy module inkeras/src/utils/module_utils.py
for dynamic import support.Input Signature and Export Utilities:
get_input_signature
to clarify behavior for different model types and ensure correct input signature construction for export. [1] [2]_infer_input_signature_from_model
to handle flexible batch dimensions and ensure compatibility with downstream exporters, always returning a flat list of input specs.